#include "RgProjScorer.hh"

#include "G4Alpha.hh"
#include "G4Step.hh"
#include "G4Track.hh"
#include "G4UnitsTable.hh"
#include "MyMacros.hh"
//#include "G4SystemOfUnits.hh"
// #include "G4ParticleDefinition.hh"

RgProjScorer::RgProjScorer(G4String name, G4int depth) : G4VPrimitiveScorer(name, depth), _fDepth(depth)
{
    _fAlpha = G4Alpha::Alpha();
    // 声明我们将会产生 3 个 map
    // 名字分别叫  <primitiveName>/CSDA  <primitiveName>/Proj  <primitiveName>/Detour
    // 这些名字在 macro 里会被用到
}

// ----------- 以下 boilerplate 直接抄即可 -----------
void RgProjScorer::Initialize(G4HCofThisEvent* HCE)
{
    MYG4TID(tid);
    // DINFOLN("[tid {} Rg/Initialize]", tid);
    //auto* const dctt = this->GetMultiFunctionalDetector();
    const auto& dcttName = this->GetMultiFunctionalDetector()->GetName();
    const auto& primName = this->GetName();
    // detector name, collection name

    // _fCSDAMap = new G4THitsMap<G4double>(dcttName, primName);
    _fProjMap = new G4THitsMap<G4double>(dcttName, primName);
    if (_fHCID < 0)
    {
        // _fHCID = HCE->GetNumberOfCollections();
        /// arg0 没有用到; 内部会调用 G4SDManager::GetSDMpointer()->GetCollectionID(detector->GetName() + "/" + primitiveName);
        //
        _fHCID = GetCollectionID(0);
    }
    // DINFOLN("[tid {} Rg/Initialize] HCID= {}, detector= {}, scorer= {}", tid, _fHCID, dcttName, primName);

    HCE->AddHitsCollection(_fHCID, _fProjMap);

    _fAccum.clear();
}

G4bool RgProjScorer::ProcessHits(G4Step* step, G4TouchableHistory*)
{
    MYG4TID(tid);
    auto* const track = step->GetTrack();
    auto* const partDef = track->GetParticleDefinition();
    const auto& ptclName = partDef->GetParticleName();

    // DINFOLN("[RgProjScorer/ProcessHits] ProcessHits()");

    // ---- 1. 过滤：只记录主粒子，且只记录 α（可改成任何粒子） ----
    if (track->GetParentID() != 0) return true;  // 只盯 primary
    // if (partDef != _fAlpha) return true;
    if (ptclName != "alpha") return true;

    // DINFOLN("[RgProjScorer/ProcessHits] stat projection range");

    const G4int tkID = track->GetTrackID();
    // 第一次见这条 track 就新开一个 slot
    if (tkID >= (G4int)_fPerTrack.size()) _fPerTrack.resize(tkID + 1);

    TrackData& d = _fPerTrack[tkID];

    // ---- 2. 第一次 step 记下起点 ----
    if (track->GetCurrentStepNumber() == 1)
    {
        d.start = track->GetPosition();
        DINFOLN("[RgProjScorer/ProcessHits] d.start= ({},{},{})", d.start.x(), d.start.y(), d.start.z());
    }

    // ---- 4. 粒子停掉或离开世界体积 → 写结果 ----
    G4bool stopped = (track->GetTrackStatus() == fStopAndKill) || (track->GetNextVolume() == nullptr);

    if (stopped)
    {
        d.end = track->GetPosition();

        // 投影轴取 Z（可改成 User 指定的 axis）
        const G4double proj = std::abs(d.end.z() - d.start.z());

        // 5. 填 hits map（index=0 代表"整个 mesh 的唯一 bin"）
        // 对于 1,1,150, 就是 z 方向的 seq
        const G4int idx = GetIndex(step);  // voxel 索引
        // DINFOLN("[tid {} Rg/ProcessHits]index {}", tid, idx);
        _fAccum[idx] = proj;

        // _fProjMap->add(idx, proj);
    }
    return true;
}

void RgProjScorer::EndOfEvent(G4HCofThisEvent*)
{
    /// 写入 EvtMap：输出 project range (mm)
    for (const auto& it : _fAccum)
    {
        //const G4int idx = it.first;
        //const auto val = it.second;
        if (it.second)
        {
            // unit: mm
            _fProjMap->add(it.first, it.second);
        }
    }
    _fAccum.clear();
}

void RgProjScorer::PrintAll()
{
    auto print = [](G4THitsMap<G4double>* m, G4String title)
    {
        G4cout << "-------- " << title << " --------" << G4endl;
        std::map<G4int, G4double*>::iterator it;
        for (it = m->GetMap()->begin(); it != m->GetMap()->end(); ++it)
            G4cout << " idx=" << it->first << "  val=" << G4BestUnit(*(it->second), "Length") << G4endl;
    };

    print(_fProjMap, "Projected Range");
}